import torch
import torch.nn as nn
import json
import random
import numpy as np
from torch.utils.data import Dataset, Subset
from utils.util import *

class MyDataset(Dataset):
    def __init__(self, args):
        super().__init__()
        self.train_data={}
        self.valid_data={}
        self.test_data={}
        self.data={}
        self.user_num=-1
        self.item_num=-1
        self.maxlen=args.maxlen
        # 从jsonl中读取数据
        with open("data_processed.jsonl", 'r') as f:
            for line in f:  # ✅ 逐行读
                line = line.strip()
                if not line:
                    continue
                data = json.loads(line)
                user_id=data['user_id']
                self.user_num=max(self.user_num, user_id)
                movies=data['movies']
                self.item_num=max(self.item_num, max([movie['id'] for movie in movies]))
                movie_len=len(movies)
                self.data[user_id]={
                    "gender":data['gender'],
                    "age":data['age'],
                    "occupation":data['occupation'],
                    "zipcode":data['zip-code'],
                    "movies":data['movies'],
                }
                
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        print(222222)
        user_id = list(self.data.keys())[idx]
        return self.data[user_id]

def item_process(user_data, feat, type, item_num, maxlen):
    user_feat={}
    item_feat={}
    if feat:
        user_feat['gender']=cat2num(user_data['gender'], "gender")
        user_feat['occupation']=cat2num(user_data['occupation'], "occupation")
        user_feat['age']=cat2num(user_data['age'], "age")
        user_feat['user_id']=user_data['user_id']

    movies = user_data['movies']
    
    if type=='train':
        if len(movies) >= 4:
            # 取除了最后两个的序列
            seq_movies = movies[:-3]
            # 正样本是除了最后三个的序列
            pos_movies = movies[1:-2]
        else:
            seq_movies = movies[:-1]
            pos_movies = movies[1:]
    elif type=="valid":
        if len(movies) >= 4:
            # 取除了最后两个的序列
            seq_movies = movies[:-2]
            # 正样本是除了最后三个的序列
            pos_movies = movies[1:-1]
        else:
            return torch.zeros(maxlen, dtype=torch.long),torch.zeros(maxlen, dtype=torch.long),torch.zeros(maxlen, dtype=torch.long)
    elif type=="test":
        if len(movies) >= 4:
            # 取除了最后两个的序列
            seq_movies = movies[:-1]
            # 正样本是除了最后三个的序列
            pos_movies = movies[1:]
        else:
            return torch.zeros(maxlen, dtype=torch.long),torch.zeros(maxlen, dtype=torch.long),torch.zeros(maxlen, dtype=torch.long)
    else:
        raise NotImplementedError(f"Type {type} is not legal!")
    
    # 负样本随机选择
    neg_movies = []
    for _ in range(len(pos_movies)):
        neg_id = 0
        while True:
            neg_id = random.randint(1, item_num)
            if neg_id not in pos_movies:
                break
        neg_movies.append(neg_id)
    
    # 处理序列长度，不足的在左侧补0
    seq_ids = [movie['id'] for movie in seq_movies]
    pos_ids = [movie['id'] for movie in pos_movies]
    ratings = [int(movie['rating']) for movie in seq_movies]
    
    # 左侧补0到maxlen长度
    if len(seq_ids) < maxlen:
        seq_ids = [0] * (maxlen - len(seq_ids)) + seq_ids
        pos_ids = [0] * (maxlen - len(pos_ids)) + pos_ids
        ratings = [0] * (maxlen - len(ratings)) + ratings
        neg_ids = [0] * (maxlen - len(neg_movies)) + neg_movies
    else:
        # 如果超过maxlen，取最后maxlen个
        seq_ids = seq_ids[-maxlen:]
        pos_ids = pos_ids[-maxlen:]
        ratings = ratings[-maxlen:]
        neg_ids = neg_movies[-maxlen:]

    item_feat['rating']=ratings

    if feat==1:
        return torch.tensor(seq_ids, dtype=torch.long),torch.tensor(pos_ids, dtype=torch.long),torch.tensor(neg_ids, dtype=torch.long), user_feat, item_feat
    elif feat==0:
        return torch.tensor(seq_ids, dtype=torch.long),torch.tensor(pos_ids, dtype=torch.long),torch.tensor(neg_ids, dtype=torch.long), None, None
    else:
        raise NotImplementedError(f"Feat {feat} is not legal!")
        

class TrainSubset(Dataset):
    def __init__(self, args):
        super().__init__()
        self.train_data={}
        self.valid_data={}
        self.test_data={}
        self.data={}
        self.user_num=-1
        self.item_num=-1
        self.maxlen=args.maxlen
        self.feat=args.feat
        # 从jsonl中读取数据
        with open("data_processed.jsonl", 'r') as f:
            for line in f:  # ✅ 逐行读
                line = line.strip()
                if not line:
                    continue
                data = json.loads(line)
                user_id=data['user_id']
                self.user_num=max(self.user_num, user_id)
                movies=data['movies']
                self.item_num=max(self.item_num, max([movie['id'] for movie in movies]))
                movie_len=len(movies)
                self.data[user_id]={
                    "user_id": user_id,
                    "gender":data['gender'],
                    "age":data['age'],
                    "occupation":data['occupation'],
                    "zipcode":data['zip-code'],
                    "movies":data['movies'],
                }
        
    def __len__(self):
        return len(self.data)
        
    def __getitem__(self, idx):
        user_id = list(self.data.keys())[idx]
        user_data = self.data[user_id]
        return item_process(user_data, self.feat, "train", self.item_num, self.maxlen)


class ValidSubset(Dataset):
    def __init__(self, args):
        super().__init__()
        self.train_data={}
        self.valid_data={}
        self.test_data={}
        self.data={}
        self.user_num=-1
        self.item_num=-1
        self.maxlen=args.maxlen
        self.args=args
        # 从jsonl中读取数据
        with open("data_processed.jsonl", 'r') as f:
            for line in f:  # ✅ 逐行读
                line = line.strip()
                if not line:
                    continue
                data = json.loads(line)
                user_id=data['user_id']
                self.user_num=max(self.user_num, user_id)
                movies=data['movies']
                self.item_num=max(self.item_num, max([movie['id'] for movie in movies]))
                movie_len=len(movies)
                self.data[user_id]={
                    "user_id":user_id,
                    "gender":data['gender'],
                    "age":data['age'],
                    "occupation":data['occupation'],
                    "zipcode":data['zip-code'],
                    "movies":data['movies'],
                }
                
    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        user_id = list(self.data.keys())[idx]
        user_data = self.data[user_id]
        return item_process(user_data, self.args.feat, "valid", self.item_num, self.maxlen)


class TestSubset(Dataset):
    def __init__(self, args):
        super().__init__()
        self.train_data={}
        self.valid_data={}
        self.test_data={}
        self.data={}
        self.user_num=-1
        self.item_num=-1
        self.maxlen=args.maxlen
        self.args=args
        # 从jsonl中读取数据
        with open("data_processed.jsonl", 'r') as f:
            for line in f:  # ✅ 逐行读
                line = line.strip()
                if not line:
                    continue
                data = json.loads(line)
                user_id=data['user_id']
                self.user_num=max(self.user_num, user_id)
                movies=data['movies']
                self.item_num=max(self.item_num, max([movie['id'] for movie in movies]))
                movie_len=len(movies)
                self.data[user_id]={
                    "user_id":user_id,
                    "gender":data['gender'],
                    "age":data['age'],
                    "occupation":data['occupation'],
                    "zipcode":data['zip-code'],
                    "movies":data['movies'],
                }
                
    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        user_id = list(self.data.keys())[idx]
        user_data = self.data[user_id]
        return item_process(user_data, self.args.feat, "test", self.item_num, self.maxlen)

